package com.xs.langchain4j_springboot.config;

 import dev.langchain4j.data.message.ChatMessage;
 import dev.langchain4j.memory.chat.MessageWindowChatMemory;
 import dev.langchain4j.store.memory.chat.ChatMemoryStore;

 import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public class PersistentChatMemoryStore implements ChatMemoryStore {

    // 使用ConcurrentHashMap存储会话ID到消息列表的映射
    private final Map<Object, List<ChatMessage>> memoryStore = new ConcurrentHashMap<>();

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        // 关键修复：返回空列表而不是null
        return memoryStore.getOrDefault(memoryId, Collections.emptyList());
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        if (messages == null) {
            // 如果传入的消息列表为null，移除该会话
            memoryStore.remove(memoryId);
        } else {
            // 复制消息列表，避免外部修改
            memoryStore.put(memoryId, new ArrayList<>(messages));
        }
    }

    @Override
    public void deleteMessages(Object memoryId) {
        memoryStore.remove(memoryId);
    }

    // 可选：添加统计方法用于调试
    public int size() {
        return memoryStore.size();
    }

    // 可选：添加方法获取所有会话ID
    public Set<Object> getAllMemoryIds() {
        return memoryStore.keySet();
    }
}